// SPDX-FileCopyrightText: 2025 UnionTech Software Technology Co., Ltd.
// SPDX-License-Identifier: MIT

package ratelimit

import (
	"net/http"
	"sync"
	"time"
)

// Limiter 限流器
type Limiter struct {
	rate     int
	tokens   int
	lastTime time.Time
	mu       sync.Mutex
}

// New 创建新的限流器
func New(rate int) *Limiter {
	return &Limiter{
		rate:     rate,
		tokens:   rate,
		lastTime: time.Now(),
	}
}

// Middleware 限流中间件
func (l *Limiter) Middleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !l.allow() {
			http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
			return
		}
		next.ServeHTTP(w, r)
	})
}

// allow 检查是否允许请求
func (l *Limiter) allow() bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(l.lastTime)

	// 计算应该补充的令牌数
	tokensToAdd := int(elapsed.Seconds() * float64(l.rate))
	if tokensToAdd > 0 {
		l.tokens = min(l.rate, l.tokens+tokensToAdd)
		l.lastTime = now
	}

	if l.tokens > 0 {
		l.tokens--
		return true
	}

	return false
}

// min 返回两个整数中的较小值
func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}
